import json

from pydantic import BaseModel, Field
from typing import Dict, List, Optional, Any
from mcp.server.sse import SseServerTransport
from mcp.server import Server
from mcp.server.fastmcp import FastMCP
from mcp.types import Tool, TextContent
from starlette.applications import Starlette
from starlette.routing import Route, Mount
import uvicorn


class ThoughtData(BaseModel):
    thought: str = Field(..., description="当前思考的内容")
    thought_number: int = Field(..., description="当前思考的序号")
    total_thought_number: int = Field(..., description="总共的思考数量")
    is_revision: Optional[int] = Field(..., description="是否修改之前的思考（可选）")
    revises_thought: Optional[int] = Field(..., description="修改的是哪个思考的序号（可选）")
    branch_from_thought: Optional[int] = Field(..., description="分支自哪个思考序号（可选）")
    branch_id: Optional[str] = Field(..., description="分支ID（可选）")
    next_thought_need: bool = Field(..., description="是否需要下一个思考步骤")


def validate_thought_data(thought_data: ThoughtData):
    if (not thought_data.thought) or (not isinstance(thought_data.thought, str)):
        raise ValueError("thought must be a string")

    if (not thought_data.thought_number) or (not isinstance(thought_data.thought_number, int)):
        raise ValueError("thought_number must be an integer")

    if (not thought_data.total_thought_number) or (not isinstance(thought_data.total_thought_number, int)):
        raise ValueError("total_thought_number must be an integer")

    if not thought_data.next_thought_need:
        # or (not isinstance(thought_data.next_thought_need, bool)):
        def str_to_bool(s):
            return {"true": True, "false": False}.get(s.lower(), None)

        if isinstance(thought_data.next_thought_need, str):
            if str_to_bool(thought_data.next_thought_need) is None:
                raise ValueError("next_thought_need must be a boolean")
            else:
                thought_data.next_thought_need = str_to_bool(thought_data.next_thought_need)

    return thought_data


class ThinkingMcpServer:
    def __init__(self):
        self.branches: Dict[str, List[ThoughtData]] = dict()

    def process_thought_data(self, thought_data: ThoughtData) -> list[TextContent]:
        try:
            validated_thought_data = validate_thought_data(thought_data)
            if validated_thought_data.thought_number > validated_thought_data.total_thought_number:
                validated_thought_data.total_thought_number = validated_thought_data.thought_number

            # 如果是分支思考，将分支添加到分支列表中
            if (getattr(validated_thought_data, "branch_from_thought", False)) and (
                    validated_thought_data.branch_id):
                if validated_thought_data.branch_id not in self.branches:
                    self.branches[validated_thought_data.branch_id] = []
                self.branches[validated_thought_data.branch_id].append(validated_thought_data)

            result = {
                "thought_number": validated_thought_data.thought_number,
                "total_thought_number": validated_thought_data.total_thought_number,
                "next_thought_need": validated_thought_data.next_thought_need,
                "branches": list(self.branches.keys())
            }
            final_result = ""
            for item in result.keys():
                final_result += f"{item}: {result[item]}\n"
            return [TextContent(type="text", text=final_result)]
        except Exception as e:
            print(e)
            return [TextContent(type="text", text=f"方法调用失败：{e}")]


mcp = FastMCP("think_step_by_step")


@mcp.tool()
async def think_step_by_step(params: ThoughtData) -> List[TextContent]:
    """
    一个用于动态和反思性解决问题的详细思考工具。
    该工具通过灵活、可适应和演进的思考过程帮助分析问题。
    每个思考步骤都可以基于、质疑或修正先前的见解，从而深化理解。
    使用场景：
        -将复杂问题分解为多个步骤
        -留有修订余地的规划与设计
        -可能需要调整方向的分析
        -初始范围不明确的问题
        -需要多步骤解决方案的问题
        -需要跨步骤保持上下文的任务
        -需要过滤无关信息的情况
    核心特性：
        -可动态调整总思考步骤数（增减）
        -可质疑或修正先前思考
        -即使到达"终点"后仍可追加思考
        -能表达不确定性并探索替代方案
        -支持非线性思考（分支或回溯）
        -生成解决方案假设
        -基于思维链验证假设
        -循环迭代直至满意
        -最终提供正确答案
    参数说明：
    thought（思考内容,中文，不要使用unicode编码）：当前思考步骤，可包含：
        -常规分析步骤
        -对先前思考的修正
        -对先前决策的质疑
        -需要更多分析的认知
        -方法论的调整
        -假设生成
        -假设验证
    next_thought_need（需后续思考）：若需更多思考则为True（即使看似已结束）,否则为 False, 必须解析成 True或 False！
    thought_number（思考序号）：当前步骤序号（可超出初始总数）
    total_thought_number（总思考数）：当前预估所需步骤数（可调整）
    is_revision（是否修正）：标记是否为修正先前思考
    revises_thought（修正目标）：若为修正，指定被修正的思考序号
    branch_from_thought（分支起点）：若为分支思考，指定源思考序号
    branch_id（分支ID）：当前分支的唯一标识（如有）
    使用建议：
    1.从预估步骤数开始，但保持调整灵活性
    2.积极质疑或修正先前思考
    3.即使"结束"后仍可追加必要思考
    4.明确表达不确定性
    5.标记修正或分支的思考步骤
    6.过滤当前步骤的无关信息
    7.适时生成解决方案假设
    8.基于思维链验证假设
    9.循环直至获得满意解
    10.最终输出单一最优答案
    11.仅在真正完成时设置next_thought_need为False
    """
    think_server = ThinkingMcpServer()
    return think_server.process_thought_data(params)


# # 初始化服务器
# app = Server("thought")
#
#
# @app.list_tools()
# async def list_tools() -> List[Tool]:
#     """
#     列出可是使用的工具
#     返回:
#         List[Tool]: 工具列表
#     """
#     return [
#         Tool(
#             name="think_step_by_step",
#             description="""
#             一个用于动态和反思性解决问题的详细思考工具。
#             该工具通过灵活、可适应和演进的思考过程帮助分析问题。
#             每个思考步骤都可以基于、质疑或修正先前的见解，从而深化理解。
#             使用场景：
#                 -将复杂问题分解为多个步骤
#                 -留有修订余地的规划与设计
#                 -可能需要调整方向的分析
#                 -初始范围不明确的问题
#                 -需要多步骤解决方案的问题
#                 -需要跨步骤保持上下文的任务
#                 -需要过滤无关信息的情况
#             核心特性：
#                 -可动态调整总思考步骤数（增减）
#                 -可质疑或修正先前思考
#                 -即使到达"终点"后仍可追加思考
#                 -能表达不确定性并探索替代方案
#                 -支持非线性思考（分支或回溯）
#                 -生成解决方案假设
#                 -基于思维链验证假设
#                 -循环迭代直至满意
#                 -最终提供正确答案
#             参数说明：
#             thought（思考内容）：当前思考步骤，可包含：
#                 -常规分析步骤
#                 -对先前思考的修正
#                 -对先前决策的质疑
#                 -需要更多分析的认知
#                 -方法论的调整
#                 -假设生成
#                 -假设验证
#             next_thought_need（需后续思考）：若需更多思考则为True（即使看似已结束）,否则为 False
#             thought_number（思考序号）：当前步骤序号（可超出初始总数）
#             total_thought_number（总思考数）：当前预估所需步骤数（可调整）
#             is_revision（是否修正）：标记是否为修正先前思考
#             revises_thought（修正目标）：若为修正，指定被修正的思考序号
#             branch_from_thought（分支起点）：若为分支思考，指定源思考序号
#             branch_id（分支ID）：当前分支的唯一标识（如有）
#             使用建议：
#             1.从预估步骤数开始，但保持调整灵活性
#             2.积极质疑或修正先前思考
#             3.即使"结束"后仍可追加必要思考
#             4.明确表达不确定性
#             5.标记修正或分支的思考步骤
#             6.过滤当前步骤的无关信息
#             7.适时生成解决方案假设
#             8.基于思维链验证假设
#             9.循环直至获得满意解
#             10.最终输出单一最优答案
#             11.仅在真正完成时设置next_thought_need为False
#             """,
#             inputSchema={
#                 "type": "object",
#                 "properties": {
#                     "thought": {
#                         "type": "string",
#                         "description": "当前的思考步骤"
#                     },
#                     "next_thought_need": {
#                         "type": "boolean",
#                         "description": "是否需要另一个思考步骤"
#                     },
#                     "thought_number": {
#                         "type": "integer",
#                         "description": "当前思考的序号",
#                         "minimum": 1
#                     },
#                     "total_thought_number": {
#                         "type": "integer",
#                         "description": "估计所需的总思考次数",
#                         "minimum": 1
#                     },
#                     "is_revision": {
#                         "type": "boolean",
#                         "description": "是否修正之前的思考"
#                     },
#                     "revises_thought": {
#                         "type": "integer",
#                         "description": "修改的是哪个思考的序号",
#                         "minimum": 1
#                     },
#                     "branch_from_thought": {
#                         "type": "integer",
#                         "description": "分支自哪个思考序号",
#                         "minimum": 1
#                     },
#                     "branch_id": {
#                         "type": "string",
#                         "description": "分支ID"
#                     },
#                     "need_more_thought": {
#                         "type": "boolean",
#                         "description": "是否需要更多思考"
#                     }
#                 },
#                 "required": ["thought", "next_thought_need", "thought_number", "total_thought_number"]
#             }
#         )
#     ]
#
#
# @app.call_tool()
# async def call_tool(name: str, arguments: dict) -> list[TextContent]:
#     if name == "think_step_by_step":
#         return think_server.process_thought_data(arguments)
#     raise ValueError(f"Unknown tool: {name}")
#
#
# sse = SseServerTransport("/messages/")
#
#
# # Handler for SSE connections
# async def handle_sse(request):
#     async with sse.connect_sse(
#             request.scope, request.receive, request._send
#     ) as streams:
#         await app.run(streams[0], streams[1], app.create_initialization_options())
#
#
# # Create Starlette app with routes
# starlette_app = Starlette(
#     debug=True,
#     routes=[
#         Route("/sse", endpoint=handle_sse),
#         Mount("/messages/", app=sse.handle_post_message)
#     ],
# )

# if __name__ == "__main__":
#     uvicorn.run(starlette_app, host="0.0.0.0", port=5052)


# 挂载SSE服务器到ASGI服务器上
app = Starlette(
    routes=[
        Mount('/', app=mcp.sse_app()),
    ]
)

if __name__ == "__main__":
    uvicorn.run("thought_mcp_server:app", host="0.0.0.0", port=5052)
